{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g_nWetWWd_ns"
      },
      "source": [
        "##### Copyright 2024 The AI Edge Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "2pHVBk_seED1"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M7vSdG6sAIQn"
      },
      "source": [
        "# TFLite Authoring Tool"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fwc5GKHBASdc"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/lite/guide/authoring\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/guide/authoring.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/guide/authoring.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/tensorflow/tensorflow/lite/g3doc/guide/authoring.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9ee074e4"
      },
      "source": [
        "TensorFlow Lite Authoring API provides a way to maintain your `tf.function` models compatibile with TensorFlow Lite.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UaWdLA3fQDK2"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DWjLcy2CvgxH"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CmkXJRDj5hTi"
      },
      "source": [
        "## TensorFlow to TensorFlow Lite compatibility issue\n",
        "\n",
        "If you want to use your TF model on devices, you need to convert it to a TFLite model to use it from TFLite interpreter.\n",
        "During the conversion, you might encounter a compatibility error because of unsupported TensorFlow ops by the TFLite builtin op set.\n",
        "\n",
        "This is a kind of annoying issue. How can you detect it earlier like the model authoring time?\n",
        "\n",
        "Note that the following code will fail on the `converter.convert()` call.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LHKqKFm5OvyQ"
      },
      "outputs": [],
      "source": [
        "@tf.function(input_signature=[\n",
        "    tf.TensorSpec(shape=[None], dtype=tf.float32)\n",
        "])\n",
        "def f(x):\n",
        "  return tf.cosh(x)\n",
        "\n",
        "# Evaluate the tf.function\n",
        "result = f(tf.constant([0.0]))\n",
        "print (f\"result = {result}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BS5bOoD50zaU"
      },
      "outputs": [],
      "source": [
        "# Convert the tf.function\n",
        "converter = tf.lite.TFLiteConverter.from_concrete_functions(\n",
        "    [f.get_concrete_function()], f)\n",
        "try:\n",
        "  fb_model = converter.convert()\n",
        "except Exception as e:\n",
        "  print(f\"Got an exception: {e}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eLU0Y9V8g_Wk"
      },
      "source": [
        "## Simple Target Aware Authoring usage\n",
        "\n",
        "We introduced Authoring API to detect the TensorFlow Lite compatibility issue during the model authoring time.\n",
        "\n",
        "You just need to add `@tf.lite.experimental.authoring.compatible` decorator to wrap your `tf.function` model to check TFLite compatibility.\n",
        "\n",
        "After this, the compatibility will be checked automatically when you evaluate your model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zVSh6VCDhbPz"
      },
      "outputs": [],
      "source": [
        "@tf.lite.experimental.authoring.compatible\n",
        "@tf.function(input_signature=[\n",
        "    tf.TensorSpec(shape=[None], dtype=tf.float32)\n",
        "])\n",
        "def f(x):\n",
        "  return tf.cosh(x)\n",
        "\n",
        "# Evaluate the tf.function\n",
        "result = f(tf.constant([0.0]))\n",
        "print (f\"result = {result}\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZWkBEqv-eUwV"
      },
      "source": [
        "If any TensorFlow Lite compatibility issue is found, it will show `COMPATIBILITY WARNING` or `COMPATIBILITY ERROR` with the exact location of the problematic op. In this example, it shows the location of `tf.Cosh` op in your tf.function model.\n",
        "\n",
        "You can also check the compatiblity log with the `\u003cfunction_name\u003e.get_compatibility_log()` method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "irwO2qdv2RPA"
      },
      "outputs": [],
      "source": [
        "compatibility_log = '\\n'.join(f.get_compatibility_log())\n",
        "print (f\"compatibility_log = {compatibility_log}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-LTVE00CiqpS"
      },
      "source": [
        "## Raise an exception for an incompatibility\n",
        "\n",
        "You can provide an option to the `@tf.lite.experimental.authoring.compatible` decorator. The `raise_exception` option gives you an exception when you're trying to evaluate the decorated model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YfPfOJm5jST4"
      },
      "outputs": [],
      "source": [
        "@tf.lite.experimental.authoring.compatible(raise_exception=True)\n",
        "@tf.function(input_signature=[\n",
        "    tf.TensorSpec(shape=[None], dtype=tf.float32)\n",
        "])\n",
        "def f(x):\n",
        "  return tf.cosh(x)\n",
        "\n",
        "# Evaluate the tf.function\n",
        "try:\n",
        "  result = f(tf.constant([0.0]))\n",
        "  print (f\"result = {result}\")\n",
        "except Exception as e:\n",
        "  print(f\"Got an exception: {e}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WXywHrR0Xjop"
      },
      "source": [
        "## Specifying \"Select TF ops\" usage\n",
        "\n",
        "If you're already aware of [Select TF ops](https://www.tensorflow.org/lite/guide/ops_select) usage, you can tell this to the Authoring API by setting `converter_target_spec`. It's the same [tf.lite.TargetSpec](https://www.tensorflow.org/api_docs/python/tf/lite/TargetSpec) object you'll use it for [tf.lite.TFLiteConverter](https://www.tensorflow.org/api_docs/python/tf/lite/TFLiteConverter) API.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "B483OwYQYG8A"
      },
      "outputs": [],
      "source": [
        "target_spec = tf.lite.TargetSpec()\n",
        "target_spec.supported_ops = [\n",
        "    tf.lite.OpsSet.TFLITE_BUILTINS,\n",
        "    tf.lite.OpsSet.SELECT_TF_OPS,\n",
        "]\n",
        "@tf.lite.experimental.authoring.compatible(converter_target_spec=target_spec, raise_exception=True)\n",
        "@tf.function(input_signature=[\n",
        "    tf.TensorSpec(shape=[None], dtype=tf.float32)\n",
        "])\n",
        "def f(x):\n",
        "  return tf.cosh(x)\n",
        "\n",
        "# Evaluate the tf.function\n",
        "result = f(tf.constant([0.0]))\n",
        "print (f\"result = {result}\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mtept13-C6uD"
      },
      "source": [
        "## Checking GPU compatibility\n",
        "\n",
        "If you want to ensure your model is compatibile with [GPU delegate](https://www.tensorflow.org/lite/performance/gpu) of TensorFlow Lite, you can set `experimental_supported_backends` of [tf.lite.TargetSpec](https://www.tensorflow.org/api_docs/python/tf/lite/TargetSpec).\n",
        "\n",
        "The following example shows how to ensure GPU delegate compatibility of your model. Note that this model has compatibility issues since it uses a 2D tensor with tf.slice operator and unsupported tf.cosh operator. You'll see two `COMPATIBILITY WARNING` with the location information."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_DzHV3KVC0T0"
      },
      "outputs": [],
      "source": [
        "target_spec = tf.lite.TargetSpec()\n",
        "target_spec.supported_ops = [\n",
        "    tf.lite.OpsSet.TFLITE_BUILTINS,\n",
        "    tf.lite.OpsSet.SELECT_TF_OPS,\n",
        "]\n",
        "target_spec.experimental_supported_backends = [\"GPU\"]\n",
        "@tf.lite.experimental.authoring.compatible(converter_target_spec=target_spec)\n",
        "@tf.function(input_signature=[\n",
        "    tf.TensorSpec(shape=[4, 4], dtype=tf.float32)\n",
        "])\n",
        "def func(x):\n",
        "  y = tf.cosh(x)\n",
        "  return y + tf.slice(x, [1, 1], [1, 1])\n",
        "\n",
        "result = func(tf.ones(shape=(4,4), dtype=tf.float32))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JvLEtCWRvvy8"
      },
      "source": [
        "## Read more\n",
        "\n",
        "For more information, please refer to:\n",
        "- [tf.function](https://www.tensorflow.org/api_docs/python/tf/function) API doc\n",
        "- [Better performance with tf.function](https://www.tensorflow.org/guide/function)\n",
        "- [TensorFlow Lite converter](https://www.tensorflow.org/lite/models/convert)\n",
        "- [TensorFlow Lite Model Analyzer](https://www.tensorflow.org/lite/guide/model_analyzer)\n",
        "- [TensorFlow Lite GPU delegate](https://www.tensorflow.org/lite/performance/gpu)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "authoring.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
